# ====================================================================================
# 文件: preprocessing/step2_CoLaKG_item_embed.py
# 描述: [V2 修复版]
#      (修复) 统一使用对象访问 ('.') 来读取 config。
# ====================================================================================

import os
import json
import torch

# 导入共享的解析器和工具函数
from base_pre_parser import get_pre_config
from pre_utils import get_sbert_embeddings


def main():
    print("--- [Step 2] CoLaKG: 物品 SBERT 嵌入生成 ---")
    config = get_pre_config()

    # --- [ V11 修复 ] ---
    conf_data = config.data_config
    conf_pre = config.preproc_config.item_enhance
    # --- [ 修复结束 ] ---

    DATA_DIR = config.DATA_DIR

    # 1. 加载 item_text.json
    # --- [ V11 修复 ] ---
    item_text_file = os.path.join(DATA_DIR, conf_data.item_meta_file)
    # --- [ 修复结束 ] ---

    print(f"[Step 2.1] 加载物品元数据: {item_text_file}")

    if not os.path.exists(item_text_file):
        print(f"错误: 物品元数据文件未找到: {item_text_file}")
        return

    try:
        with open(item_text_file, 'r') as f:
            item_meta = json.load(f)
    except json.JSONDecodeError as e:
        print(f"错误: 解析 JSON 文件失败: {e}")
        return

    # 2. 准备文本列表
    # --- [ V11 修复 ] ---
    n_items = config.data_config.n_items
    # --- [ 修复结束 ] ---

    if n_items is None:
        print("错误: 'configs/amazon-book.yaml' 中必须定义 'data.n_items' (物品总数)")
        return

    print(f"  > 总物品数 (n_items): {n_items}")

    texts = []
    missing_items = 0
    for i in range(n_items):
        text = item_meta.get(str(i))
        if text:
            texts.append(text)
        else:
            texts.append("")  # 使用空字符串填充
            missing_items += 1

    print(f"  > 准备了 {len(texts)} 个物品的文本 (其中 {missing_items} 个为空)")

    # 3. 生成 SBERT 嵌入
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"[Step 2.2] 开始生成 SBERT 嵌入 (设备: {device})")

    # --- [ V11 修复 ] ---
    sbert_model = conf_pre.sbert_model_name
    # --- [ 修复结束 ] ---

    embeddings = get_sbert_embeddings(
        texts,
        sbert_model,
        device,
        batch_size=128
    )

    # 4. 保存嵌入
    # --- [ V11 修复 ] ---
    output_file = os.path.join(DATA_DIR, conf_pre.output_embed_file)
    # --- [ 修复结束 ] ---

    torch.save(embeddings, output_file)

    print(f"\n✅ [Step 2] 成功完成!")
    print(f"物品 SBERT 嵌入已保存到: {output_file}")


if __name__ == "__main__":
    main()